import torch
import torch.nn as nn
from 封装ST_GCN.model.conv_temporal_graphical import ConvTemporalGraphical


def no_residual(x):
    # 无残差连接时，直接返回0
    return 0


def identity_residual(x):
    # 恒等残差连接，直接返回输入x
    return x


class STGCNBlock(nn.Module):
    def __init__(self,
                 in_channels,  # 输入通道数
                 out_channels,  # 输出通道数
                 kernel_size,  # 卷积核大小
                 stride=1,  # 卷积步长，默认为1
                 dropout=0,  # 丢弃率，默认为0
                 residual=True):  # 是否使用残差连接，默认为True
        super().__init__()

        assert len(kernel_size) == 2
        assert kernel_size[0] % 2 == 1
        # 计算填充大小，这里只对时间维度进行填充
        padding = ((kernel_size[0] - 1) // 2, 0)

        # 时空图卷积层
        self.gcn = ConvTemporalGraphical(in_channels, out_channels, kernel_size[1])
        # 时间卷积层，包含批量归一化、ReLU激活函数、卷积层、批量归一化和Dropout层
        self.tcn = nn.Sequential(
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(
                out_channels,
                out_channels,
                (kernel_size[0], 1),
                (stride, 1),
                padding,
            ),
            nn.BatchNorm2d(out_channels),
            nn.Dropout(dropout, inplace=True),
        )

        if not residual:
            # 如果不使用残差连接，将residual设置为no_residual函数
            self.residual = no_residual
            # self.residual = lambda x: 0

        elif (in_channels == out_channels) and (stride == 1):
            # 如果输入通道数等于输出通道数且步长为1，将residual设置为identity_residual函数
            self.residual = identity_residual
            # self.residual = lambda x: x

        else:
            # 否则，创建包含卷积层和批量归一化层的残差连接
            self.residual = nn.Sequential(
                nn.Conv2d(
                    in_channels,
                    out_channels,
                    kernel_size=1,
                    stride=(stride, 1)),
                nn.BatchNorm2d(out_channels),
            )

        # ReLU激活函数
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x, A):
        # 计算残差
        res = self.residual(x)
        # 时空图卷积操作
        x = self.gcn(x, A)
        # 时间卷积操作，并加上残差
        x = self.tcn(x) + res

        # 应用ReLU激活函数
        return self.relu(x)
